{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compress and Evaluate Video Generation Models"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/PrunaAI/pruna/blob/v|version|/docs/tutorials/video_generation.ipynb\">\n",
    "    <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "| Component | Details |\n",
    "|-----------|---------|\n",
    "| **Goal** | Showcase a standard workflow for optimizing and evaluating a video generation model |\n",
    "| **Model** |[Wan-AI/Wan2.1-T2V-1.3B](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B) |\n",
    "| **Dataset** |  [nannullna/laion_subset](https://huggingface.co/datasets/nannullna/laion_subset) |\n",
    "| **Device** | 1 x H100 (80GB) |\n",
    "| **Optimization Algorithms** | compiler(torch_compile), kernel(flash_attn3) |\n",
    "| **Evaluation Metrics** | `total time`, `latency`, `througput`, `co2_emissions`, and `energy_consumed` |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting Started\n",
    "\n",
    "To install the required dependencies, you can run the following command:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install pruna\n",
    "%pip install ftfy imageio imageio-ffmpeg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For more information about how to install Pruna, please refer to the [Installation](https://docs.pruna.ai/en/stable/setup/install.html) page.\n",
    "\n",
    "Then, we will set the device to the best available option to maximize the optimization process's benefits. However, in this case, we recommend using a GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Load the Model\n",
    "\n",
    "First, we must load the original model using the diffusers library to ensure it fits into memory. In this example, we will use a light model compatible with most of the consumer-grade GPUs, [Wan-AI/Wan2.1-T2V-1.3B](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B).\n",
    "\n",
    "Pruna works at least as well with larger models, like the model version of Wan 2.1 14B or HuyuanVideo. The choice to use a smaller model is simply because it’s a good starting point, so feel free to use any [text-to-video model available on Hugging Face](https://huggingface.co/models?pipeline_tag=text-to-video&sort=trending)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers import AutoencoderKLWan, WanPipeline\n",
    "\n",
    "model_id = \"Wan-AI/Wan2.1-T2V-1.3B-Diffusers\"\n",
    "\n",
    "vae = AutoencoderKLWan.from_pretrained(model_id, subfolder=\"vae\", torch_dtype=torch.float32)\n",
    "\n",
    "pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we have loaded the pipeline, we can run some inference and check the output. The standard prompt structure for a video is **Subject + Subject Action + Scene**, which can become more complex as we add descriptions and details like the lighting, point of view, or visual style to achieve specific and refined results.\n",
    "\n",
    "Remember that you can improve the quality of the video by increasing the number of frames, the number of inference steps, and the guidance scale, but this will also increment the time and amount of resources required to generate the video."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers.utils import export_to_video\n",
    "\n",
    "prompt = \"A dog runs on the beach, realistic.\"\n",
    "negative_prompt = \"Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards\"  # noqa: E501\n",
    "\n",
    "with torch.no_grad():\n",
    "    output = pipe(\n",
    "        prompt=prompt,\n",
    "        negative_prompt=negative_prompt,\n",
    "        height=480,\n",
    "        width=832,\n",
    "        num_frames=33,\n",
    "        guidance_scale=3.0,\n",
    "        num_inference_steps=15,\n",
    "        generator=torch.Generator(device=device).manual_seed(42),\n",
    "    ).frames[0]\n",
    "\n",
    "export_to_video(output, \"base_video.mp4\", fps=15)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "<video controls width=\"100%\" style=\"max-width:640px; height:auto;\">\n",
    "    <source src=\"../assets/videos/video_generation/base_video.mp4\" type=\"video/mp4\">\n",
    "    Your browser does not support HTML5 video.\n",
    "</video>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, the model has generated a nice short video based on our prompt."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Define the SmashConfig\n",
    "\n",
    "Now that we have correctly loaded and tested our base model, let's continue by defining the `SmashConfig` to customize the optimizations we want to apply when smashing.\n",
    "\n",
    "Take into account that not all optimization algorithms are available for all models, so you can learn about the requirements and compatibility in the [Algorithms Overview](https://docs.pruna.ai/en/stable/compression.html).\n",
    "\n",
    "In the current optimization, we will use [torch_compile](https://docs.pruna.ai/en/stable/compression.html#torch-compile) to make it more efficient and [flash_attn3](https://docs.pruna.ai/en/stable/compression.html#flash-attn3) to speed up the model.\n",
    "\n",
    "Let's define the `SmashConfig` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pruna import SmashConfig\n",
    "\n",
    "smash_config = SmashConfig({\n",
    "    \"torch_compile\": {\"target\": \"module_list\"},\n",
    "    \"flash_attn3\": {},\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Smash the Model\n",
    "\n",
    "Next, we need to apply our defined `SmashConfig` by smashing our model. The `smash` function will be in charge of this, so we just need to pass the `model` and the `smash_config`. To evaluate and compare the models in the upcoming sections, we will make a deep copy of the base model.\n",
    "\n",
    "Time to smash! This will take around 20 seconds, depending on the configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "\n",
    "from pruna import smash\n",
    "\n",
    "copy_pipe = copy.deepcopy(pipe).to(\"cpu\")\n",
    "smashed_pipe = smash(\n",
    "    model=pipe,\n",
    "    smash_config=smash_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we will have an optimized smashed model, so let's check how it works using the previous prompt.\n",
    "\n",
    "Consider that if you are using `torch_compile` as a compiler, you can expect the first inference warmup to take longer than the actual inference.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    output = smashed_pipe(\n",
    "        prompt=prompt,\n",
    "        negative_prompt=negative_prompt,\n",
    "        height=480,\n",
    "        width=832,\n",
    "        num_frames=33,\n",
    "        guidance_scale=3.0,\n",
    "        num_inference_steps=15,\n",
    "        generator=torch.Generator(device=device).manual_seed(42),\n",
    "    ).frames[0]\n",
    "\n",
    "export_to_video(output, \"smashed_video.mp4\", fps=15)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "<video controls width=\"100%\" style=\"max-width:640px; height:auto;\">\n",
    "    <source src=\"../assets/videos/video_generation/smashed_video.mp4\" type=\"video/mp4\">\n",
    "    Your browser does not support HTML5 video.\n",
    "</video>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can observe, it has also generated a short video similar to the original model.\n",
    "\n",
    "If you notice a significant difference, it might be due to the model, the configuration, the hardware, etc. We encourage you to retry the optimization process or try out different configurations and models to find the best fit for your use case. However, feel free to reach out to us on [Discord](https://discord.gg/JFQmtFKCjd) if you have any questions or feedback."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Evaluate the Smashed Model\n",
    "\n",
    "Now that we have our smashed model, the key question is how much has improved with our optimization. For this, we can run an evaluation of the performance using the `EvaluationAgent`. In this case, we will include metrics like the `total_time`, `latency`, `throughput`, `co2_emissions`, and `energy_consumed`.\n",
    "\n",
    "A complete list of the available metrics can be found in [Evaluation](https://docs.pruna.ai/en/stable/reference/evaluation.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pruna import PrunaModel\n",
    "from pruna.data.pruna_datamodule import PrunaDataModule\n",
    "from pruna.evaluation.evaluation_agent import EvaluationAgent\n",
    "from pruna.evaluation.metrics import (\n",
    "    CO2EmissionsMetric,\n",
    "    EnergyConsumedMetric,\n",
    "    LatencyMetric,\n",
    "    ThroughputMetric,\n",
    "    TotalTimeMetric,\n",
    ")\n",
    "from pruna.evaluation.task import Task\n",
    "\n",
    "# Define the metrics. Increment the number of iterations and\n",
    "# warmup iterations to get a more accurate result.\n",
    "metrics = [\n",
    "    TotalTimeMetric(n_iterations=3, n_warmup_iterations=1),\n",
    "    LatencyMetric(n_iterations=3, n_warmup_iterations=1),\n",
    "    ThroughputMetric(n_iterations=3, n_warmup_iterations=1),\n",
    "    CO2EmissionsMetric(n_iterations=3, n_warmup_iterations=1),\n",
    "    EnergyConsumedMetric(n_iterations=3, n_warmup_iterations=1),\n",
    "]\n",
    "\n",
    "# Define the datamodule\n",
    "datamodule = PrunaDataModule.from_string(\"LAION256\")\n",
    "datamodule.limit_datasets(10)\n",
    "\n",
    "# Define the task and the evaluation agent\n",
    "task = Task(metrics, datamodule=datamodule, device=device)\n",
    "eval_agent = EvaluationAgent(task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate the smashed model and offload it to CPU\n",
    "from pruna.engine.utils import move_to_device\n",
    "\n",
    "move_to_device(smashed_pipe, device)\n",
    "smashed_model_results = eval_agent.evaluate(smashed_pipe)\n",
    "move_to_device(smashed_pipe, \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate the base model and offload it to CPU\n",
    "base_pipe = PrunaModel(model=copy_pipe)\n",
    "move_to_device(base_pipe, device)\n",
    "base_model_results = eval_agent.evaluate(base_pipe)\n",
    "move_to_device(base_pipe, \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualize and compare the evaluation results of the base and smashed models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "| Metric | Base Model | Smashed Model | Improvement % |\n",
       "|-----|-----|-----|-----|\n",
       "| total_time | 460992.1875000 ms | 265793.1718750 ms | 42.34% |\n",
       "| latency | 153664.0625000 ms/num_iterations | 88597.7239583 ms/num_iterations | 42.34% |\n",
       "| throughput | 0.0000065 num_iterations/ms | 0.0000113 num_iterations/ms | 73.44% |\n",
       "| co2_emissions | 0.0031181 kgCO2e | 0.0018072 kgCO2e | 42.04% |\n",
       "| energy_consumed | 0.0556424 kWh | 0.0322483 kWh | 42.04% |"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Markdown, display  # noqa\n",
    "\n",
    "\n",
    "def make_comparison_table(base_model_results, smashed_model_results):  # noqa\n",
    "    header = \"| Metric | Base Model | Smashed Model | Improvement % |\\n\"\n",
    "    header += \"|\" + \"-----|\" * 4 + \"\\n\"\n",
    "    rows = []\n",
    "\n",
    "    for base, smashed in zip(base_model_results, smashed_model_results):\n",
    "        base_result = base.result\n",
    "        smashed_result = smashed.result\n",
    "        if base.higher_is_better:\n",
    "            diff = ((smashed_result - base_result) / base_result) * 100\n",
    "        else:\n",
    "            diff = ((base_result - smashed_result) / base_result) * 100\n",
    "        row = f\"| {base.name} | {base_result:.7f} {base.metric_units or ''}\"\n",
    "        row += f\"| {smashed_result:.7f} {smashed.metric_units or ''} | {diff:.2f}% |\"\n",
    "        rows.append(row)\n",
    "    return header + \"\\n\".join(rows)\n",
    "\n",
    "\n",
    "display(Markdown(make_comparison_table(base_model_results, smashed_model_results)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, the smashed model is much more efficient than the base model. It runs almost 2x faster both overall and per iteration and has a higher throughput. Moreover, the energy consumption and CO₂ emissions were also reduced, meaning that the compressed model is not only faster but also more environmentally friendly, consuming less electricity and producing less carbon footprint. This results are consistent with the expected ones as both the compiler and the kernel are designed to improve the performance of the model.\n",
    "\n",
    "So, we can save the optimized model to disk or share it with others. Note that some optimizations, such as `torch_compile`, are device dependent and will be re-applied when loading the model on a different device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the model to disk\n",
    "smashed_pipe.save_pretrained(\"Wan2.1-T2V-1.3B-smashed\")\n",
    "# Load the model from disk\n",
    "# smashed_pipe = PrunaModel.from_pretrained(\"Wan2.1-T2V-1.3B-smashed/\")\n",
    "\n",
    "# Save the model to HuggingFace\n",
    "# smashed_pipe.push_to_hub(\"PrunaAI/Wan2.1-T2V-1.3B-smashed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusions\n",
    "\n",
    "In this tutorial, we have gone over the standard workflow for optimizing and evaluating a text-to-video model.\n",
    "\n",
    "We started loading the base model and defining the SmashConfig with the desired optimization algorithms and parameters. Then we smashed the base model, obtaining an optimized version, and we ensured the improvement in performance by running an evaluation with the EvaluationAgent.\n",
    "\n",
    "The results show that we can significantly increase the inference speed andreduce the energy consumption, while maintaining a high level of output quality. This makes it easy to explore trade-offs and iterate on configurations to find the best optimization strategy for your specific use case.\n",
    "\n",
    "Check out our other [tutorials](https://docs.pruna.ai/en/stable/docs_pruna/tutorials/index.html) for more examples on how to optimize and evaluate image generation models or LLM models."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pruna-tutorials",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
